import platgo as pg
import numpy as np

"""
------------------------------- Reference -----------------------------
 K. Deb, A. Pratap, S. Agarwal, and T. Meyarivan, A fast and elitist
 multiobjective genetic algorithm: NSGA-II, IEEE Transactions on
 Evolutionary Computation, 2002, 6(2): 182-197.
"""


class NSGA2(pg.Algorithm):

    type: dict = {'single': False, 'multi': True, 'many': False, 'real': True, 'binary': True, 'permutation': True,
                  "large": False, 'expensive': False, 'constrained': True, 'preference': False, 'multimodal': False,
                  'sparse': False, 'gradient': False}

    def __init__(self, maxgen: int, problem: pg.Problem) -> None:
        super().__init__(maxgen=maxgen, problem=problem)
        self.name = "NSGA2"
        self.xov = pg.operators.XovSbx()  # 模拟二进制交叉
        self.mut = pg.operators.MutPol(problem)  # 多项式变异

    def go(self, N: int = None, population: pg.Population = None):
        assert N or population, "N and population can't be both None"
        if population is None:
            pop = self.problem.init_pop(N)
        else:
            pop = population
            self.problem.N = pop.decs.shape[0]
        self.problem.cal_obj(pop)
        _, frontno, cd = environmental_selection(pop, pop.N, self.problem)  # 第一次用_,接收pop是因为初始种群保留

        while self.not_terminal(pop):
            # TODO 未考虑编码形式
            matingpool = pg.utils.tournament_selection(2, pop.N, frontno, -cd)
            p1, p2 = pg.utils.random_selection(pop[matingpool])
            offspring = self.xov(pop, p1, p2)
            offspring = self.mut(offspring)
            self.problem.cal_obj(offspring)  # 计算子代种群目标函数值
            temp_pop = offspring + pop  # 合并种群
            pop, frontno, cd = environmental_selection(temp_pop, pop.N, self.problem)
        return pop


def environmental_selection(pop: pg.Population, N: [int, float], problem):
    """
    The environmental selection of NSGA-II
    """
    # nd_sort
    if problem.type["constrained"]:
        frontno, maxfront = pg.utils.nd_sort(pop.objv, pop.cv, N)
    else:
        frontno, maxfront = pg.utils.nd_sort(pop.objv, N)
    next = frontno < maxfront
    # calculate the crowding distance
    cd = pg.utils.crowding_distance(pop, frontno)

    # select the soltions in the last front based on their crowding distances
    last = np.argwhere(frontno == maxfront)
    rank = np.argsort(-cd[last], axis=0)
    next[last[rank[0:(N - np.sum(next))]]] = True  # 先计算next（也就是不是最大前沿上的）的总数，让N-sum即为需要在最后一前沿需要挑选的解
    # 此时选取最后前沿上rank（即依据cd选取的）最好的解，再使其的next变为Ture
    # pop for next generation
    pop = pop[next]  # 选取next（即选中的解）中的解
    frontno = frontno[next]  # 选取这些解的前沿
    cd = cd[next]
    return pop, frontno, cd
